﻿/*
===========================================================================
Immemorsion LGPL Source code
Copyright trszdev - rahimyanov@gmail.com

This file is part of Immemorsion.

Immemorsion is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

Immemorsion is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with Immemorsion.  If not, see <http://www.gnu.org/licenses/>.
===========================================================================
*/
using System;
using System.Collections.Generic;
using System.Linq;
using System.Diagnostics;
using System.IO;
using System.Text;
using Immemorsion.PlatformApi;

namespace Immemorsion
{
    public class ProcessMemoryScanner : WindowsProcessCommunicator, IProcessMemoryScanner
    {
        public ProcessMemoryScanner(Process process)
            : base(process)
        {}

        private static readonly WindowsAllocationProtect ReadFlag =
            WindowsAllocationProtect.PAGE_READWRITE;

        private static WindowsSystemInfo GetSystemInfo()
        {
            var result = new WindowsSystemInfo();
            WindowsApi.GetSystemInfo(out result);
            return result;
        }

        private WindowsMemoryBasicInformation MakeVirtualQuery(long baseAddress)
        {
            var result = new WindowsMemoryBasicInformation();
            WindowsApi.VirtualQueryEx(ProcessHandle, (IntPtr)baseAddress, out result, 28);
            return result;
        }

        private static KeyValuePair<byte, int> GetFirstByteInfo(IByteMemoryPattern memoryPattern)
        {
            var first = memoryPattern
                .Select((a, index) => new { Index = index, Byte = a.Byte, Mask = a.ByteMask })
                .First(a => a.Mask == ByteMask.Exact);
            var result = new KeyValuePair<byte, int>(first.Byte, first.Index);
            return result;
        }


        // Return type changed to list due to low perfomance of yield
        private List<MemoryAddress> _SearchPattern(IByteMemoryPattern memoryPattern)
        {
            var systemInfo = GetSystemInfo();
            var proc_min_address = (long)systemInfo.minimumApplicationAddress;
            var proc_max_address = (long)systemInfo.maximumApplicationAddress;
            var result = new List<MemoryAddress>();
            var firstByteInfo = GetFirstByteInfo(memoryPattern);
            var firstByte = firstByteInfo.Key;
            var patternOffset = firstByteInfo.Value;

            
            var memHandler = new ProcessMemoryHandle(ProcessHandle);
            var slice = new byte[memoryPattern.Count];

            while (proc_min_address < proc_max_address)
            {
                var memoryInformation = MakeVirtualQuery(proc_min_address);
                if (memoryInformation.State != WindowsMemoryState.MEM_COMMIT ||
                    memoryInformation.Protect != (int)ReadFlag)
                    goto NextPage;

                memHandler.BaseAddress = memoryInformation.BaseAddress;
                memHandler.Length = memoryInformation.RegionSize;
                memHandler.Update();
                if (memHandler.IsBroken)
                    goto NextPage;
                var buffer = memHandler.Bytes;

                for (var i = Array.IndexOf<byte>(buffer, firstByte, patternOffset) - patternOffset;
                i > -1;
                i = Array.IndexOf<byte>(buffer, firstByte, i + patternOffset + 1) - patternOffset)
                {
                    if (buffer.Length - i < memoryPattern.Count) break;
                    Array.Copy(buffer, i, slice, 0, memoryPattern.Count);
                    if (memoryPattern.Suits(slice))
                    {
                        result.Add(new MemoryAddress(memoryInformation.BaseAddress + i));
                    }
                }
                NextPage:
                proc_min_address += memoryInformation.RegionSize;
            }
            return result;
        }

        //TODO: Finish if pattern don't have exact bytes
        public IEnumerable<MemoryAddress> SearchPattern(IByteMemoryPattern memoryPattern, IMemoryRegion memoryRegion)
        {
            var hasExactByte = memoryPattern.Any(a => a.ByteMask == ByteMask.Exact);
            if (hasExactByte)
            {
                var enumerator = _SearchPattern(memoryPattern)
                    .Where(a => a.Value >= memoryRegion.Start && a.Value <= memoryRegion.End);
                foreach (var result in enumerator)
                    yield return result;
            }
            else
            {
                //Finish here
                var lowerByte = memoryPattern
                    .Where(a => a.ByteMask == ByteMask.Lower)
                    .Select((a, index) => new { Byte = a.Byte, ByteMask = a.ByteMask, Index = index })
                    .Aggregate((a, b) => a.Byte < b.Byte ? a : b);
                var higherByte = memoryPattern
                    .Where(a => a.ByteMask == ByteMask.Higher)
                    .Select((a, index) => new { Byte=a.Byte,ByteMask=a.ByteMask,Index=index })
                    .Aggregate((a, b) => a.Byte > b.Byte ? a : b);
                var preferredByte = 255 - higherByte.Byte > lowerByte.Byte ? lowerByte : higherByte;
                if (preferredByte.ByteMask==ByteMask.Higher)
                {
                   
                }
                else
                {
                    
                }
            }   
            yield break;
        }

        public IEnumerable<MemoryAddress> SearchPattern(IByteMemoryPattern memoryPattern)
        {
            foreach (var result in _SearchPattern(memoryPattern))
                yield return result;
            yield break;
        }
    }
}
